import tensorflow as tf
import os
import numpy as np
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

VER = 'v1.0'
VECTOR_DIR = '_save/x30_extract_feature_and_save.py/v2.0/vectors'
ALPHA = 0.001
ITERS = 200
FILE_NAME = os.path.basename(__file__)
MODEL_SAVE_DIR = os.path.join('_save', FILE_NAME, 'model')
os.makedirs(MODEL_SAVE_DIR, exist_ok=True)
MODEL_SAVE_PATH = os.path.join(MODEL_SAVE_DIR, 'ckpt')

print('Reading data')
x = []
y = []
vec_len = None
yi = 0
idx2label, label2idx = {}, {}
for sub_dir in os.listdir(VECTOR_DIR):
    sub_dir_path = os.path.join(VECTOR_DIR, sub_dir)
    idx2label[yi] = sub_dir
    label2idx[sub_dir] = yi
    for file in os.listdir(sub_dir_path):
        file_path = os.path.join(sub_dir_path, file)
        vec = np.loadtxt(file_path).reshape(-1)
        if vec_len is None:
            vec_len = len(vec)
            print(f'vec len:', vec_len)
        x.append(vec)
        y.append(yi)
    yi += 1
n_cls = len(idx2label)
print('n_cls', n_cls)
x = np.float32(x)
y = np.int64(y)
print('x', x.shape)
print('y', y.shape)

x_train, x_val, y_train, y_val = train_test_split(x, y, train_size=0.9, random_state=1, shuffle=True)
print('x_train', x_train.shape, x_train.dtype)
print('x_val', x_val.shape, x_val.dtype)
print('y_train', y_train.shape, y_train.dtype)
print('y_val', y_val.shape, y_val.dtype)

ph_x = tf.placeholder(tf.float32, (None, vec_len), 'ph_x')
ph_y = tf.placeholder(tf.int64, (None, ), 'ph_y')
pred = tf.layers.Dense(n_cls, activation=None)(ph_x)
loss = tf.reduce_mean(
    tf.nn.sparse_softmax_cross_entropy_with_logits(labels=ph_y, logits=pred)
)
accuracy = tf.reduce_mean(
    tf.cast(
        tf.equal(ph_y, tf.argmax(pred, axis=1)),
        tf.float32
    )
)
optim = tf.train.AdamOptimizer(learning_rate=ALPHA).minimize(loss)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    loss_his, acc_his = [], []
    group = int(np.ceil(ITERS / 20))
    for i in range(ITERS):
        _, lossv, accv = sess.run([optim, loss, accuracy], feed_dict={ph_x: x_train, ph_y: y_train})
        loss_his.append(lossv)
        acc_his.append(accv)
        if i % group == 0 or i == ITERS - 1:
            print(f'#{i + 1}: loss = {lossv}, acc = {accv}')

    print('Saving model ...')
    saver = tf.train.Saver(tf.global_variables(), max_to_keep=6)
    path_prefix = saver.save(sess, MODEL_SAVE_PATH, global_step=0)
    print(f'Saved to {path_prefix}*')

    lossv, accv = sess.run([loss, accuracy], feed_dict={ph_x: x_val, ph_y: y_val})
    print(f'Test: loss = {lossv}, acc = {accv}')

    spr = 1
    spc = 2
    spn = 0
    plt.figure(figsize=[12, 6])

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('loss')
    plt.plot(loss_his)
    plt.grid()

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('acc')
    plt.plot(acc_his)
    plt.grid()

    plt.show()
